#!/usr/bin/python
# vim: set fileencoding=utf8 :

import glob,sys,os,re,codecs
from pylab import *
import tables
from PIL import Image
import ImageFont,ImageDraw
from scipy.ndimage import filters,measurements
from ocrolib import ligatures,docproc
from ocrolib.ocroio import renumber_labels
from llpy.sutils import *
import traceback
from scipy.misc import imsave

import argparse
parser = argparse.ArgumentParser(description = "Generate text line training data")
parser.add_argument('-o','--base',default='linegen')
parser.add_argument('-f','--fontlist',default=None)
parser.add_argument('-F','--fontdir',default="/home/tmb/bin/Fonts/")
parser.add_argument('-l','--lines',default=None)
parser.add_argument('-m','--maxlines',default=30,help='max # lines for each directory')
parser.add_argument('-e','--degradations',default=None)
parser.add_argument('-s','--sizes',default="36,48")
parser.add_argument('-k','--kind',default="ligatures",help='characters to generate')
parser.add_argument('-d','--display',action="store_true")
parser.add_argument('-D','--debug_show',default=None,help="select a class for stepping through")
parser.add_argument('-N','--nsamples',default=999999999,help="max number of samples",type=int)
args = parser.parse_args()

sizes = [int(x) for x in args.sizes.split(",")]

fonts = """AmericanTypewriter.ttf AmericanTypewriterBold.ttf
AmericanTypewriterCondensed.ttf AmericanTypewriterCondensedBold.ttf
AmericanTypewriterCondensedLight.ttf AmericanTypewriterLight.ttf
Chalkboard.ttf ChalkboardBold.ttf
GenevaCY.ttf Georgia.ttf GeorgiaBold.ttf
GeorgiaBoldItalic.ttf GeorgiaItalic.ttf GillSans.ttf GillSansBold.ttf
GillSansBoldItalic.ttf GillSansItalic.ttf GillSansLight.ttf
GillSansLightItalic.ttf HelveticaNeue.ttf HelveticaNeueBold.ttf
HelveticaNeueBoldItalic.ttf HelveticaNeueCondensedBlack.ttf
HelveticaNeueCondensedBold.ttf HelveticaNeueItalic.ttf
HelveticaNeueLight.ttf HelveticaNeueLightItalic.ttf
OptimaBold.ttf OptimaBoldItalic.ttf OptimaExtraBlack.ttf OptimaItalic.ttf
OptimaRegular.ttf Palatino.ttf PalatinoBold.ttf PalatinoBoldItalic.ttf
PalatinoItalic.ttf SkiaRegular.ttf StoneSansITCTTBold.ttf
StoneSansSemITCTTSemi.ttf StoneSansSemITCTTSemiIta.ttf Vera.ttf VeraBI.ttf
VeraBd.ttf VeraIt.ttf VeraMoBI.ttf VeraMoBd.ttf VeraMoIt.ttf VeraMono.ttf
VeraSe.ttf VeraSeBd.ttf arial.ttf arialbd.ttf arialbi.ttf
ariali.ttf ariblk.ttf comic.ttf comicbd.ttf cour.ttf courbd.ttf courbi.ttf
couri.ttf framd.ttf framdit.ttf georgia.ttf georgiab.ttf georgiai.ttf
georgiaz.ttf impact.ttf kartika.ttf l_10646.ttf lucon.ttf micross.ttf
pala.ttf palab.ttf palabi.ttf palai.ttf tahoma.ttf tahomabd.ttf times.ttf
timesbd.ttf timesbi.ttf timesi.ttf trebuc.ttf trebucbd.ttf trebucbi.ttf
trebucit.ttf verdana.ttf verdanab.ttf verdanai.ttf verdanaz.ttf
vrinda.ttf"""

################################################################

fonts = re.sub(" ","\n",fonts)

if args.fontlist is not None:
    with open(args.fontlist) as stream:
        fonts = stream.read()

fontlist = []
for font in fonts.split("\n"):
    font = font.strip()
    if font=="": continue
    if font[0]!="/": font = args.fontdir + font
    fontlist.append(font)

# sigma +/-   threshold +/-
deglist = """
0.5 0.0   0.5 0.0
1.0 0.0   0.4 0.0
1.0 0.0   0.5 0.0
1.0 0.0   0.6 0.0
"""

if args.degradations is not None:
    with open(args.degradations) as stream:
        deglist = stream.read()

degradations = []
for deg in deglist.split("\n"):
    deg = deg.strip()
    if deg=="": continue
    deg = [float(x) for x in deg.split()]
    degradations.append(deg)

def rgeometry(image,eps=0.03,delta=0.3):
    m = array([[1+eps*randn(),0.0],[eps*randn(),1.0+eps*randn()]])
    w,h = image.shape
    c = array([w/2.0,h/2])
    d = c-dot(m,c)+array([randn()*delta,randn()*delta])
    return interpolation.affine_transform(image,m,offset=d,order=1,mode='constant',cval=image[0,0])

if args.debug_show:
    ion(); gray()

def bounding_box(a):
    a = array(a>0,'i')
    l = measurements.find_objects(a)
    if len(l)<1: return (0,0,0,0)
    ys,xs = l[0]
    # y0,x0,y1,x1
    return (ys.start,xs.start,ys.stop,xs.stop)

lines = """\
When, in the course of human events, it becomes necessary for one
people to dissolve the political bands which have connected them with
another, and to assume among the powers of the earth, the separate and
equal station to which the laws of nature and of nature's God entitle
them, a decent respect to the opinions of mankind requires that they
should declare the causes which impel them to the separation.  We hold
these truths to be self-evident, that all men are created equal, that
they are endowed by their Creator with certain unalienable rights,
that among these are life, liberty and the pursuit of happiness. That
to secure these rights, governments are instituted among men, deriving
their just powers from the consent of the governed. That whenever any
form of government becomes destructive to these ends, it is the right
of the people to alter or to abolish it, and to institute new
government, laying its foundation on such principles and organizing
its powers in such form, as to them shall seem most likely to effect
their safety and happiness. Prudence, indeed, will dictate that
governments long established should not be changed for light and
transient causes; and accordingly all experience hath shown that
mankind are more disposed to suffer, while evils are sufferable, than
to right themselves by abolishing the forms to which they are
accustomed. But when a long train of abuses and usurpations, pursuing
invariably the same object evinces a design to reduce them under
absolute despotism, it is their right, it is their duty, to throw off
such government, and to provide new guards for their future
security. --Such has been the patient sufferance of these colonies;
and such is now the necessity which constrains them to alter their
former systems of government. The history of the present King of Great
Britain is a history of repeated injuries and usurpations, all having
in direct object the establishment of an absolute tyranny over these
states. To prove this, let facts be submitted to a candid world.
""".split("\n")

if args.lines:
    with codecs.open(args.lines,encoding="utf-8") as stream:
        lines = stream.read().split("\n")

lines = [line for line in lines if len(line)>=2]
print "got",len(lines),"lines"

base = args.base
print "base",base
os.system("rm -rf "+base)
os.mkdir(base)

def crop(image,pad=1):
    [[r,c]] = measurements.find_objects(array(image==0,'i'))
    r0 = r.start
    r1 = r.stop
    c0 = c.start
    c1 = c.stop
    image = image[r0-pad:r1+pad,c0-pad:c1+pad]
    return image

last_font = None
last_fontfile = None

def genline(text,fontfile=args.fontdir+"/Chalkboard.ttf",size=36,sigma=0.5,threshold=0.5):
    global image,draw,last_font,last_fontfile
    if last_fontfile!=fontfile:
        last_font = ImageFont.truetype(fontfile,size)
        last_fontfile = fontfile
    font = last_font
    image = Image.new("L",(6000,200))
    draw = ImageDraw.Draw(image)
    draw.rectangle((0,0,6000,6000),fill="white")
    draw.text((20,20),text,fill="black",font=font)
    a = asarray(image,'f')
    a = a*1.0/amax(a)
    if sigma>0.0:
        a = filters.gaussian_filter(a,sigma)
    a += clip(randn(*a.shape)*0.2,-0.25,0.25)
    a = rgeometry(a)
    a = array(a>threshold,'f')
    a = crop(a,pad=3)
    # FIXME add grid warping here
    # clf(); ion(); gray(); imshow(a); ginput(1,0.1)
    del draw
    del image
    return a

pageno = 1

jobs = []
for fontfile in fontlist:
    for size in sizes:
        for (sigma,ssigma,threshold,sthreshold) in degradations:
            sigma += (2*rand()-1)*ssigma
            threshold += (2*rand()-1)*sthreshold
            jobs.append((fontfile,size,sigma,threshold))


random.shuffle(jobs)

print "#jobs (pages)",len(jobs)

for fontfile,size,sigma,threshold in jobs:
    pagedir = "%s/%04d"%(base,pageno)
    os.mkdir(pagedir)
    print "===",pagedir,fontfile,size,sigma,threshold
    with open(pagedir+".info","w") as stream:
        stream.write("%s\n"%fontfile)
        stream.write("%d %f %f\n"%(size,sigma,threshold))
    lineno = 1
    sel = range(len(lines))
    random.shuffle(sel)
    for i in sel:
        line = lines[i]
        line = line.strip()
        if len(line)<1: continue
        try:
            image = genline(text=line,fontfile=fontfile,size=size,sigma=sigma,threshold=threshold)
        except:
            traceback.print_exc()
            continue
        if amin(image.shape)<10: continue
        if amax(image)<0.5: continue
        if args.display:
            clf(); imshow(image); ginput(1,0.1)
        fname = "%s/%04d/01%04d"%(base,pageno,lineno)
        imsave(fname+".png",image)
        imsave(fname+".bin.png",image)
        with open(fname+".gt.txt","w") as stream:
            stream.write(line+"\n")
        lineno += 1
        if lineno>=args.maxlines: break
    pageno += 1
